/*
* Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
* Description: Hccl One-Sided Service
* Author: yanzhenghang
* Create: 2024-09-06
*/

#include "hccl_one_sided_service.h"
#include "device_capacity.h"

namespace hccl {
using namespace std;

HcclOneSidedService::HcclOneSidedService(unique_ptr<HcclSocketManager> &socketManager,
    unique_ptr<NotifyPool> &notifyPool)
    : IHcclOneSidedService(socketManager, notifyPool)
{
}

HcclResult HcclOneSidedService::IsUsedRdma(RankId remoteRankId, bool &useRdma)
{
    DevType deviceType;
    CHK_RET(hrtGetDeviceType(deviceType));

    RankInfo_t localRankInfo = (rankTable_->rankList).at(localRankInfo_.userRank);
    RankInfo_t remoteRankInfo = (rankTable_->rankList).at(remoteRankId);
    if (deviceType == DevType::DEV_TYPE_910B) {
        // 外部使能RDMA，或者节点间通信
        if (GetExternalInputIntraRoceSwitch() || localRankInfo.serverId != remoteRankInfo.serverId) {
            useRdma = true;
            return HCCL_SUCCESS;
        }

        // 同一节点的 PCIe 连接判断
        s32 localDeviceId = localRankInfo_.devicePhyId;
        s32 remoteDeviceId = remoteRankInfo.deviceInfo.devicePhyId;
        LinkTypeInServer linkType = LinkTypeInServer::RESERVED_LINK_TYPE;
        CHK_RET(hrtGetPairDeviceLinkType(static_cast<u32>(localDeviceId), static_cast<u32>(remoteDeviceId), linkType));
        if (linkType != LinkTypeInServer::HCCS_TYPE) {
            HCCL_ERROR("[HcclOneSidedService][IsUsedRdma]localDeviceId: %d, remoteDeviceId: %d, linkType %u is not supported",
                localDeviceId, remoteDeviceId, linkType);
            return HCCL_E_NOT_SUPPORT;
        }

        // 节点内通信，默认不使用 RDMA
        useRdma = false;
        return HCCL_SUCCESS;
    } else if (deviceType == DevType::DEV_TYPE_910_93) {
        if (GetExternalInputIntraRoceSwitch() || localRankInfo.superPodId != remoteRankInfo.superPodId) {
            useRdma = true;
            return HCCL_SUCCESS;
        }

        useRdma = false;
        return HCCL_SUCCESS;
    }

    // 其他情况默认使用 RDMA
    useRdma = true;
    return HCCL_SUCCESS;
}

void HcclOneSidedService::RegMemRdma(MemInfo &memInfo, RankId remoteRankId, HcclMemDesc &localMemDesc,
    bool &isRegistedSucc)
{
    auto addr = memInfo.addr;
    auto size = memInfo.size;
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(addr), size);
    shared_ptr<LocalRdmaRmaBuffer> localbufferPtr = make_shared<LocalRdmaRmaBuffer>(netDevCtx_, addr, size,
        memInfo.type);
    auto resultPair = localRdmaRmaBufferMgr_.Add(tempKey, localbufferPtr);
    if (resultPair.first == localRdmaRmaBufferMgr_.End()) {
        // 输入key是表中某一个最相近key的交集、子集。返回空迭代器
        HCCL_ERROR("[HcclOneSidedService][RegMemRdma]The memory that is expected to be"\
            " registered overlaps with the memory that has been registered, please check params");
        throw logic_error("[HcclOneSidedService][RegMemRdma]memory register failed.");
    }
    if (resultPair.second) {
        EXCEPTION_THROW_IF_ERR(localbufferPtr->Init(), "[HcclOneSidedService][RegMemRdma]localbuffer init failed.");
    }

    // 已注册：输入key是表中某一最相近key的全集。 返回添加该key的迭代器，及false
    // 未注册：输入key是表中某一最相近key的空集。 返回添加成功的迭代器，及true
    std::string logInfo = resultPair.second ? "Register memory success!"
                        : "Memory is already registed, just increase the reference count.";
    // false：内存已注册过，仅引用计数增加
    isRegistedSucc = resultPair.second;
    HCCL_INFO("[HcclOneSidedService][RegMemRdma]:%s Add key {%p, %llu}", logInfo.c_str(), addr, size);

    std::string tempLocalMemDesc = resultPair.first->second.buffer->Serialize();
    if (tempLocalMemDesc.empty()) {
        HCCL_ERROR("[HcclOneSidedService][RegMemRdma]tempLocalMemDesc is empty.");
        throw logic_error("[HcclOneSidedService][RegMemRdma]tempLocalMemDesc is empty.");
    }

    MemDescData* ptr = reinterpret_cast<MemDescData*>(localMemDesc.desc);
    ptr->localRankId = localRankInfo_.userRank;
    ptr->remoteRankId = remoteRankId;
    ptr->CopyFromString(tempLocalMemDesc);

    std::string tempDesc = ptr->CopyToString();
    desc2LocalRdmaRmaBufferMap_.emplace(tempDesc, resultPair.first->second.buffer);
    HCCL_DEBUG("[HcclOneSidedService][RegMemRdma]localRankId: %u, remoteRankId: %u, length: %u %u, MapSize: %u",
        ptr->localRankId, remoteRankId, tempLocalMemDesc.length(), tempDesc.length(),
        desc2LocalRdmaRmaBufferMap_.size());
}

void HcclOneSidedService::RegMemIpc(MemInfo &memInfo, RankId remoteRankId, HcclMemDesc &localMemDesc,
    bool &isRegistedSucc)
{
    auto addr = memInfo.addr;
    auto size = memInfo.size;
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(addr), size);
    shared_ptr<LocalIpcRmaBuffer> localbufferPtr = make_shared<LocalIpcRmaBuffer>(netDevCtx_, addr, size, memInfo.type);
    auto resultPair = localIpcRmaBufferMgr_.Add(tempKey, localbufferPtr);
    if (resultPair.first == localIpcRmaBufferMgr_.End()) {
        // 输入key是表中某一个最相近key的交集、子集。返回空迭代器
        HCCL_ERROR("[HcclOneSidedService][RegMemIpc]The memory that is expected to be"\
            " registered overlaps with the memory that has been registered, please check params");
        throw logic_error("[HcclOneSidedService][RegMemIpc]memory register failed.");
    }
    if (resultPair.second) {
        EXCEPTION_THROW_IF_ERR(localbufferPtr->Init(), "[HcclOneSidedService][RegMemIpc]localbuffer init failed.");
    }

    // 已注册：输入key是表中某一最相近key的全集。 返回添加该key的迭代器，及false
    // 未注册：输入key是表中某一最相近key的空集。 返回添加成功的迭代器，及true
    std::string logInfo = resultPair.second ? "Register memory success!"
                        : "Memory is already registed, just increase the reference count.";
    HCCL_INFO("[HcclOneSidedService][RegMemIpc]:%s Add key {%p, %llu}", logInfo.c_str(), addr, size);
    // false：内存已注册过，仅引用计数增加
    isRegistedSucc = resultPair.second;
    std::string tempLocalMemDesc = resultPair.first->second.buffer->Serialize();
    if (tempLocalMemDesc.empty()) {
        HCCL_ERROR("[HcclOneSidedService][RegMemIpc]tempLocalMemDesc is empty.");
        throw logic_error("[HcclOneSidedService][RegMemIpc]tempLocalMemDesc is empty.");
    }

    MemDescData* ptr = reinterpret_cast<MemDescData*>(localMemDesc.desc);
    ptr->localRankId = localRankInfo_.userRank;
    ptr->remoteRankId = remoteRankId;
    ptr->CopyFromString(tempLocalMemDesc);

    std::string tempDesc = ptr->CopyToString();
    desc2LocalIpcRmaBufferMap_.emplace(tempDesc, resultPair.first->second.buffer);
    HCCL_DEBUG("[HcclOneSidedService][RegMemIpc]localRankId: %u, remoteRankId: %u, length: %u %u, MapSize: %u",
        ptr->localRankId, remoteRankId, tempLocalMemDesc.length(), tempDesc.length(),
        desc2LocalIpcRmaBufferMap_.size());
}

HcclResult HcclOneSidedService::RegMem(void* addr, u64 size, HcclMemType type, RankId remoteRankId,
    HcclMemDesc &localMemDesc)
{
    constexpr u32 MAX_REGISTED_MEM = 256;
    if (registedMemCnt_ >= MAX_REGISTED_MEM) {
        HCCL_ERROR("[HcclOneSidedService][RegMem]The number of registered memory "\
            "exceeds the upper limit[%u]", MAX_REGISTED_MEM);
        return HCCL_E_UNAVAIL;
    }

    if (isUsedRdmaMap_.find(remoteRankId) == isUsedRdmaMap_.end()) {
        bool useRdma = true;
        CHK_RET(IsUsedRdma(remoteRankId, useRdma));
        isUsedRdmaMap_[remoteRankId] = useRdma;
    }
    bool isRegistedSucc = false;
    MemInfo memInfo(addr, size, type);
    if (isUsedRdmaMap_[remoteRankId]) {
        RegMemRdma(memInfo, remoteRankId, localMemDesc, isRegistedSucc);
    } else {
        RegMemIpc(memInfo, remoteRankId, localMemDesc, isRegistedSucc);
    }
    if (isRegistedSucc) {
        registedMemCnt_++;
    }
    return HCCL_SUCCESS;
}

void HcclOneSidedService::DeregMemRdma(const HcclMemDesc &localMemDesc, bool &isDeregistedSucc)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(localMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    HCCL_DEBUG("[HcclOneSidedService][DeregMemRdma]tempDesc: %s, length: %u, mapSize: %u",
        tempDesc.c_str(), tempDesc.length(), desc2LocalRdmaRmaBufferMap_.size());
    auto iter = desc2LocalRdmaRmaBufferMap_.find(tempDesc);
    if (iter == desc2LocalRdmaRmaBufferMap_.end()) {
        HCCL_ERROR("[HcclOneSidedService][DeregMemRdma]memory is not registered, please register first");
        throw logic_error("memory is not registered");
    }
    std::shared_ptr<LocalRdmaRmaBuffer> tempLocalBuffer = iter->second;
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(tempLocalBuffer->GetAddr()), tempLocalBuffer->GetSize());
    if (localRdmaRmaBufferMgr_.Del(tempKey)) {
        // 删除成功：输入key是表中某一最相近key的全集，计数-1后为0，返回true
        desc2LocalRdmaRmaBufferMap_.erase(tempDesc);
        isDeregistedSucc = true;
        HCCL_INFO("[HcclOneSidedService][DeregMemRdma]Memory reference count is 0, deregister memory.");
    } else {
        // 删除失败：输入key是表中某一最相近key的全集，计数不为0（存在其他remoteRank使用），返回false
        HCCL_INFO("[HcclOneSidedService][DeregMemRdma]Memory reference count is larger than 0"
            "(used by other RemoteRank), do not deregister memory.");
    }
}

void HcclOneSidedService::DeregMemIpc(const HcclMemDesc &localMemDesc, bool &isDeregistedSucc)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(localMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    HCCL_DEBUG("[HcclOneSidedService][DeregMemIpc]tempDesc: %s, length: %u, mapSize: %u",
        tempDesc.c_str(), tempDesc.length(), desc2LocalIpcRmaBufferMap_.size());
    auto iter = desc2LocalIpcRmaBufferMap_.find(tempDesc);
    if (iter == desc2LocalIpcRmaBufferMap_.end()) {
        HCCL_ERROR("[HcclOneSidedService][DeregMemIpc]memory is not registered, please register first");
        throw logic_error("memory is not registered");
    }
    std::shared_ptr<LocalIpcRmaBuffer> tempLocalBuffer = iter->second;
    BufferKey<uintptr_t, u64> tempKey(reinterpret_cast<uintptr_t>(tempLocalBuffer->GetAddr()), tempLocalBuffer->GetSize());
    if (localIpcRmaBufferMgr_.Del(tempKey)) {
        // 删除成功：输入key是表中某一最相近key的全集，计数-1后为0，返回true
        desc2LocalIpcRmaBufferMap_.erase(tempDesc);
        isDeregistedSucc = true;
        HCCL_INFO("[HcclOneSidedService][DeregMemIpc]Memory reference count is 0, deregister memory.");
    } else {
        // 删除失败：输入key是表中某一最相近key的全集，计数不为0（存在其他remoteRank使用），返回false
        HCCL_INFO("[HcclOneSidedService][DeregMemIpc]Memory reference count is larger than 0 "
            "(used by other RemoteRank), do not deregister memory.");
    }
}

HcclResult HcclOneSidedService::DeregMem(const HcclMemDesc &localMemDesc)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(localMemDesc.desc);
    u32 remoteRankId = ptr->remoteRankId;
    if (registedMemCnt_ == 0) {
        HCCL_ERROR("[HcclOneSidedService][DeregMem]The number of registered memory is 0, "\
            "please register first.");
        return HCCL_E_NOT_FOUND;
    }

    if (isUsedRdmaMap_.find(remoteRankId) == isUsedRdmaMap_.end()) {
        bool useRdma = true;
        CHK_RET(IsUsedRdma(remoteRankId, useRdma));
        isUsedRdmaMap_[remoteRankId] = useRdma;
    }
    bool isDeregistedSucc = false;
    if (isUsedRdmaMap_[remoteRankId]) {
        DeregMemRdma(localMemDesc, isDeregistedSucc);
    } else {
        DeregMemIpc(localMemDesc, isDeregistedSucc);
    }
    if (isDeregistedSucc) {
        registedMemCnt_--;
    }
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedService::Grant(const HcclMemDesc &localMemDesc, const ProcessInfo &remoteProcess)
{
    // 查找localIpcRmaBuffer对象
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(localMemDesc.desc);
    std::string tempDesc = ptr->CopyToString();

    HCCL_DEBUG("[HcclOneSidedService][Grant]tempDesc: %s, length: %u, mapSize: %u",
        tempDesc.c_str(), tempDesc.length(), desc2LocalIpcRmaBufferMap_.size());
    auto iter = desc2LocalIpcRmaBufferMap_.find(tempDesc);
    if (iter == desc2LocalIpcRmaBufferMap_.end()) {
        HCCL_ERROR("[HcclOneSidedService][Grant]memory is not registered, please register first");
        throw logic_error("memory is not registered");
    }
    std::shared_ptr<LocalIpcRmaBuffer> tempLocalBuffer = iter->second;

    // 设置白名单
    EXCEPTION_THROW_IF_ERR(tempLocalBuffer->Grant(remoteProcess.pid, remoteProcess.sdid), "[HcclOneSidedService][Grant]Grant error");
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedService::SetupRemoteRankInfo(RankId remoteRankId, HcclRankLinkInfo &remoteRankInfo)
{
    // 检查 rankId 是否有效
    CHK_PRT_RET(rankTable_->rankList.size() <= remoteRankId,
        HCCL_ERROR("[HcclOneSidedService][SetupRemoteRankInfo] the size of rankList is less than remoteRankId[%u].",
            remoteRankId), HCCL_E_NOT_FOUND);

    RankInfo_t tempRankInfo = rankTable_->rankList.at(remoteRankId);
    remoteRankInfo.userRank = tempRankInfo.rankId;
    remoteRankInfo.devicePhyId = tempRankInfo.deviceInfo.devicePhyId;

    // 检查 deviceIp 是否为空
    CHK_PRT_RET(tempRankInfo.deviceInfo.deviceIp.empty(),
        HCCL_ERROR("[HcclOneSidedService][SetupRemoteRankInfo] deviceIp is empty. RemoteRankId is [%u]",
            remoteRankId), HCCL_E_NOT_FOUND);
    remoteRankInfo.ip = tempRankInfo.deviceInfo.deviceIp[0];

    if (isUsedRdmaMap_.find(remoteRankId) != isUsedRdmaMap_.end() && !isUsedRdmaMap_[remoteRankId]) {
        bool useSuperPodMode = false;
        CHK_RET(IsSuperPodMode(useSuperPodMode));

        HcclIpAddress localVnicIp = HcclIpAddress(localRankInfo_.devicePhyId);
        HcclIpAddress remoteVnicIp = HcclIpAddress(remoteRankInfo.devicePhyId);
        RankInfo_t tRankInfo = rankTable_->rankList.at(localRankInfo_.userRank);

        if (useSuperPodMode) {
            CHK_RET(hrtRaGetSingleSocketVnicIpInfo(localRankInfo_.devicePhyId, DeviceIdType::DEVICE_ID_TYPE_SDID,
                tRankInfo.superDeviceId, localVnicIp));
            CHK_RET(hrtRaGetSingleSocketVnicIpInfo(localRankInfo_.devicePhyId, DeviceIdType::DEVICE_ID_TYPE_SDID,
                tempRankInfo.superDeviceId, remoteVnicIp));
        } else {
            CHK_RET(hrtRaGetSingleSocketVnicIpInfo(localRankInfo_.devicePhyId, DeviceIdType::DEVICE_ID_TYPE_PHY_ID,
                localRankInfo_.devicePhyId, localVnicIp));
            CHK_RET(hrtRaGetSingleSocketVnicIpInfo(localRankInfo_.devicePhyId, DeviceIdType::DEVICE_ID_TYPE_PHY_ID,
                remoteRankInfo.devicePhyId, remoteVnicIp));
        }

        localRankInfo_.ip = localVnicIp;
        remoteRankInfo.ip = remoteVnicIp;
    }
    remoteRankInfo.port = tempRankInfo.deviceInfo.port == 0 || tempRankInfo.deviceInfo.port == HCCL_INVALID_PORT ?
        HETEROG_CCL_PORT : tempRankInfo.deviceInfo.port;
    remoteRankInfo.socketsPerLink = 1;
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedService::CreateConnection(RankId remoteRankId, const HcclRankLinkInfo &remoteRankInfo,
    std::shared_ptr<HcclOneSidedConn> &tempConn, const std::string &commIdentifier, s32 timeoutSec)
{
    EXECEPTION_CATCH(tempConn = std::make_shared<HcclOneSidedConn>(netDevCtx_, localRankInfo_, remoteRankInfo,
        socketManager_, localRdmaRmaBufferMgr_, notifyPool_, localIpcRmaBufferMgr_, dispatcher_, commIdentifier,
        isUsedRdmaMap_[remoteRankId]), return HCCL_E_PTR);
    CHK_SMART_PTR_NULL(tempConn);
    CHK_RET(tempConn->Connect(timeoutSec));
    return HCCL_SUCCESS;
}

HcclResult HcclOneSidedService::ExchangeMemDesc(RankId remoteRankId, const HcclMemDescs &localMemDescs,
    HcclMemDescs &remoteMemDescs, u32 &actualNumOfRemote, const std::string &commIdentifier, s32 timeoutSec)
{
    std::shared_ptr<HcclOneSidedConn> tempConn;

    // 检查连接是否已存在，不存在则创建连接，已存在则直接获取
    auto it = oneSidedConns_.find(remoteRankId);
    if (it == oneSidedConns_.end()) {
        HcclRankLinkInfo remoteRankInfo;
        CHK_RET(SetupRemoteRankInfo(remoteRankId, remoteRankInfo));
        CHK_RET(CreateConnection(remoteRankId, remoteRankInfo, tempConn, commIdentifier, timeoutSec));
        oneSidedConns_.emplace(remoteRankId, tempConn);
    } else {
        tempConn = it->second;
    }

    // HCCS下进行权限授予
    if (!isUsedRdmaMap_[remoteRankId]) {
        u32 pid;
        SalGetBareTgid(reinterpret_cast<u32*>(&pid));
        RankId localRankId = localRankInfo_.userRank;
        u32 sdid = rankTable_->rankList.at(localRankId).superDeviceId;
        u32 serverId = rankTable_->rankList.at(localRankId).serverIdx;

        // 收发进程信息
        ProcessInfo localProcess = {pid, sdid, serverId};
        ProcessInfo remoteProcess = {0};
        CHK_RET(tempConn->ExchangeIpcProcessInfo(localProcess, remoteProcess));
        remoteProcess.sdid = localProcess.serverId == remoteProcess.serverId ? INVALID_INT : remoteProcess.sdid;

        for (u32 i = 0; i < localMemDescs.arrayLength; i++) {
            CHK_RET(Grant(localMemDescs.array[i], remoteProcess));
        }
    }

    // 交换内存描述符
    CHK_RET(tempConn->ExchangeMemDesc(localMemDescs, remoteMemDescs, actualNumOfRemote));
    return HCCL_SUCCESS;
}

void HcclOneSidedService::EnableMemAccess(const HcclMemDesc &remoteMemDesc, HcclMem &remoteMem)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    u32 remoteRank = ptr->localRankId;
    if (oneSidedConns_.find(remoteRank) == oneSidedConns_.end()) {
        HCCL_ERROR("[HcclOneSidedService][EnableMemAccess]connection not found, remoteRank[%u], "\
            "please exchange mem desc to create connection first.", remoteRank);
        throw logic_error("[HcclOneSidedService][EnableMemAccess]connection not found.");
    }
    oneSidedConns_.at(remoteRank)->EnableMemAccess(remoteMemDesc, remoteMem);
}

void HcclOneSidedService::DisableMemAccess(const HcclMemDesc &remoteMemDesc)
{
    const MemDescData* ptr = reinterpret_cast<const MemDescData*>(remoteMemDesc.desc);
    u32 remoteRank = ptr->localRankId;
    if (oneSidedConns_.find(remoteRank) == oneSidedConns_.end()) {
        HCCL_ERROR("[HcclOneSidedService][DisableMemAccess]connection not found by remoteRankId[%u], "\
            "please exchange mem desc to create connection first.", remoteRank);
        throw logic_error("[HcclOneSidedService][DisableMemAccess]connection not found.");
    }
    oneSidedConns_.at(remoteRank)->DisableMemAccess(remoteMemDesc);
}

void HcclOneSidedService::BatchPut(RankId remoteRankId, const HcclOneSideOpDesc* desc, u32 descNum,
    const rtStream_t &stream)
{
    auto it = oneSidedConns_.find(remoteRankId);
    if (it == oneSidedConns_.end()) {
        HCCL_ERROR("[HcclMemCommunication][BatchPut] Cann't find oneSidedConn by remoteRank %u", remoteRankId);
        throw out_of_range("Cann't find oneSidedConn by remoteRank.");
    }
    it->second->BatchWrite(desc, descNum, stream);
}

void HcclOneSidedService::BatchGet(RankId remoteRankId, const HcclOneSideOpDesc* desc, u32 descNum,
    const rtStream_t &stream)
{
    auto it = oneSidedConns_.find(remoteRankId);
    if (it == oneSidedConns_.end()) {
        HCCL_ERROR("[HcclMemCommunication][BatchGet] Cann't find oneSidedConn by remoteRank %u", remoteRankId);
        throw out_of_range("Cann't find oneSidedConn by remoteRank.");
    }
    it->second->BatchRead(desc, descNum, stream);
}
}